package uebung03_w;

import java.awt.BasicStroke;
import java.awt.BorderLayout;
import java.awt.Color;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.Paint;
import java.awt.RenderingHints;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.event.MouseMotionAdapter;
import java.awt.event.MouseWheelEvent;
import java.awt.event.MouseWheelListener;
import java.awt.geom.AffineTransform;
import java.awt.geom.NoninvertibleTransformException;
import java.awt.geom.Point2D;
import java.awt.image.BufferedImage;
import java.util.ArrayList;

import javax.imageio.ImageIO;
import javax.swing.BoxLayout;
import javax.swing.JComponent;
import javax.swing.JFrame;
import javax.swing.JPanel;

import org.ejml.simple.*; // Matrix Library

import tools.Arrow;
import tools.ColorConversions;
import tools.Derivative;
import tools.Gauss;


/*
 * Kanade-Lucas Optical Flow
 * 
 * 
 */

/*
 * References
 * [1]	Jean-Yves Bouget (Intel Corp. Microprocessor Research Labs) 
 * 		Pyramidal Implementation of the Lucas Kanade Feature Tracker - Description of the algorithm
 * 		http://robots.stanford.edu/cs223b04/algo_tracking.pdf
 * 		§2.4 summary of the pyramidal tracking algorithm
 *  
 */


public class IterativeKLT {
	static int bilerp(int[] pix, double x, double y, int w, int h) {
		if (x>w || y>h || x<0 || y<0)
			return 0;
		double down = Math.min(h-1, Math.ceil(y)), up = Math.floor(y), left = Math.floor(x), right = Math.min(w-1,Math.ceil(x));
		try {
			double i1 = pix[(int)(left+up*w)], i2 = pix[(int)(right+up*w)], i3 = pix[(int)(left+down*w)], i4 = pix[(int)(right+down*w)];
			double x1 = (1-(x-left))*i1+(x-left)*i2, x2 = (1-(x-left))*i3+(x-left)*i4;
			int val = (int)(x2*(y-up)+x1*(1-(y-up)));
//			System.out.println(val);
//			System.out.println(val);
			return val;
			
		}
		catch (Exception e) {
			e.printStackTrace();
			System.out.println(up+" "+down+" "+left+" "+right+" "+(y-up)+" "+(x-left));
		}
		return 0;
	}
	public static void main(String[] args) {
		try {
			final BufferedImage imgI = ImageIO.read(IterativeKLT.class.getClassLoader().getResourceAsStream("res/lena1.png"));
			final BufferedImage imgJ = ImageIO.read(IterativeKLT.class.getClassLoader().getResourceAsStream("res/lena2.png"));
			
			final int w = imgI.getWidth(), h = imgI.getHeight();
			
			final BufferedImage bi = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
			
			int[] rgbI = imgI.getRGB(0, 0, w, h, new int[w*h], 0, w);
			int[] rgbJ = imgJ.getRGB(0, 0, w, h, new int[w*h], 0, w);
			
			// helligkeitsbild
			int[] I = tools.ColorConversions.getFastLuminance(rgbI, new int[w*h]);
			int[] J = tools.ColorConversions.getFastLuminance(rgbJ, new int[w*h]);
			
			double sigma = 1.6;
			
//			int[] gaussIntensities = Gauss.gaussXY(I, w, h, sigma, new int[w*h]);
//			System.arraycopy(gaussIntensities, 0, I, 0, I.length);
//
//			gaussIntensities = Gauss.gaussXY(J, w, h, 3, gaussIntensities);
//			System.arraycopy(gaussIntensities, 0, J, 0, J.length);
			
			int[] pixels = ColorConversions.buildRGB(I, J, new int[w*h], new int[w*h]);
			bi.setRGB(0, 0, w, h, pixels, 0, w);
			
			//Ableitungen
			final int[] dx = Derivative.dx(I, w, h, new int[w*h]);
			final int[] dy = Derivative.dy(I, w, h, new int[w*h]);
			
			final int[] dt = Derivative.dt(J, I, new int[w*h]);

			final int neighbours = 10, step = 3;
			final ArrayList<Point2D> uv = new ArrayList<Point2D>();
			final ArrayList<Point2D> xy = new ArrayList<Point2D>();
			final ArrayList<Point2D> eigenvalues = new ArrayList<Point2D>();
			
			final double limit = 10;
//			final double norm = ((2*neighbours)+1)*((2*neighbours)+1);
			final int Z = 5;
			for (double y = neighbours, endy = h-neighbours; y<endy; y+=step) {
				for (double x = neighbours, endx = w-neighbours; x<endx; x+=step) {
					
					double u = 0, v = 0, eig1 = 0, eig2 = 0;
					for (int z = 0; z<Z; z++) {
						double xx = 0, yy = 0, yx = 0, tx = 0, ty = 0, weight = 0;
						for (int j = -neighbours; j<neighbours; j++) {
							for (int i = -neighbours; i<neighbours; i++) {
								int index = (int)(x+i+(y+j)*w), index2 = Math.min(Math.max((int)(x+u+i+(y+v+j)*w), dx.length), 0);
								double gauss = Gauss.gaussian1d(Math.sqrt(i*i+j*j), neighbours/3.);
//								gauss = 1;
//								int dxv = bilerp(dx, x, y, w, h), dyv = bilerp(dy, x, y, w, h), dtv = bilerp(dt, x+u, y+v, w, h);
								int dxv = (int)subpixel(dx, x, y, w), dyv = (int)subpixel(dy, x, y, w), dtv = (int)subpixel(dt, x+u, y+v, w);
//								int dxv = dx[index], dyv = dy[index], dtv = dt[index2];
//								System.out.println(dxv+" "+dyv+" "+dtv);
								xx += dxv*dxv*gauss;
								yy += dyv*dyv*gauss;
								yx += dyv*dxv*gauss;
								tx -= dtv*dxv*gauss;
								ty -= dtv*dyv*gauss;
								weight += gauss;
							}
						}
						
						xx /= weight;
						yy /= weight;
						yx /= weight;
						tx /= weight;
						ty /= weight;
						
						if (Double.compare(0, xx*yy-yx*yx)!=0) {
							SimpleMatrix sm = new SimpleMatrix(2, 2);
							sm.set(0, 0, xx);
							sm.set(1, 0, yx);
							sm.set(0, 1, yx);
							sm.set(1, 1, yy);
							SimpleEVD<SimpleMatrix> evd = sm.eig();
//							System.out.println(evd.getEigenvalue(0)+" "+evd.getEigenvalue(1));
							if (evd.getEigenvalue(0).real>limit || evd.getEigenvalue(1).real>limit) {
								
								double u_ = ((tx*yy-yx*ty)/(xx*yy-yx*yx));
								double v_ = ((ty-yx*u_)/yy);
								u+=u_;
								v+=v_;
//								if (u*u+v*v<norm) {
//									u *= Math.sqrt(norm);
//									v *= Math.sqrt(norm);
								eig1 = evd.getEigenvalue(0).real;
								eig2 = evd.getEigenvalue(1).real;

//								}
							}
						} 
						else {
//							System.out.println("verkackt");
						}

					}
					eigenvalues.add(new Point2D.Double(eig1, eig2));
					uv.add(new Point2D.Double(x+u,y+v));
					xy.add(new Point2D.Double(x,y));

				}
			}
			double sumeig1 = 0, sumeig2 = 0;
			double maxd = 0;
			for (Point2D p : eigenvalues) {
				sumeig1 += p.getX();
				sumeig2 += p.getY();
				maxd = Math.max(maxd, Math.max(p.getX(), p.getY()));
			}
			
			sumeig1 /= (double)eigenvalues.size();
 			sumeig2 /= (double)eigenvalues.size();
 			final double max = maxd*0.1;
 			
 			final double averageig1 = sumeig1, averageig2 = sumeig2;
//			// build pyramid representations of I and J
//			for (int i = 0;i<4;i++){
//				
//			}
			final AffineTransform zoom = new AffineTransform();
			final JComponent something = new JComponent() {
				{setPreferredSize(new Dimension(imgI.getWidth(),imgI.getHeight()));}
				Arrow arrow = new Arrow();
				
				@Override
				public void paint(Graphics g) {
					Graphics2D g2 = (Graphics2D)g;
					g2.transform(zoom);
					g2.drawImage(bi, 0, 0, this);
					g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
					g2.setRenderingHint(RenderingHints.KEY_STROKE_CONTROL, RenderingHints.VALUE_STROKE_PURE);
					g2.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
					g2.setStroke(new BasicStroke(0.1f));
					for (int i = 0; i<xy.size(); i++) {
						Point2D p = eigenvalues.get(i);
//						if (p.getX()>averageig1 && p.getY()>averageig2) {
						if (Math.max(p.getX(), p.getY())>max) {
							System.out.println(xy.get(i).distance(uv.get(i)));
							arrow.setArrow(xy.get(i), uv.get(i), 1, 5);
							g2.setColor(Color.magenta);
							g2.fill(arrow);
							g2.setColor(Color.white);
							g2.draw(arrow);
							
						}
					}
					// pfeile malen
					super.paint(g);
				}
			};
			something.addMouseWheelListener(new MouseWheelListener() {

				@Override
				public void mouseWheelMoved(MouseWheelEvent e) {

					double scale = Math.pow(1.01, e.getWheelRotation());

					try {
						Point2D p = new Point2D.Double(e.getX(), e.getY());
						zoom.inverseTransform(p, p);
						zoom.concatenate(AffineTransform.getTranslateInstance(p.getX(), p.getY()));
						zoom.concatenate(AffineTransform.getScaleInstance(scale,scale));
						zoom.concatenate(AffineTransform.getTranslateInstance(-p.getX(), -p.getY()));

					} catch (NoninvertibleTransformException nte) {
						nte.printStackTrace();
					}
					something.repaint();
				}
			});
			final int scale = 10, nn = 2*neighbours+1;
			final BufferedImage window = new BufferedImage(nn, nn, BufferedImage.TYPE_INT_ARGB);
			final BufferedImage movedwindow = new BufferedImage(nn, nn, BufferedImage.TYPE_INT_ARGB);
			
			final JComponent windowShow = new JComponent() {
				{setPreferredSize(new Dimension(scale*(nn), scale*(nn)));}
				public void paint(Graphics g) {
					g.drawImage(window.getScaledInstance(scale*(nn), scale*(nn), BufferedImage.SCALE_FAST), 0, 0, this);
				};
			};
			final JComponent movedWindowShow = new JComponent() {
				{setPreferredSize(new Dimension(scale*(nn), scale*(nn)));}
				public void paint(Graphics g) {
					g.drawImage(movedwindow.getScaledInstance(scale*(nn), scale*(nn), BufferedImage.SCALE_FAST), 0, 0, this);
				};
			};
			
			something.addMouseListener(new MouseAdapter() {
				
				@Override
				public void mouseClicked(MouseEvent event) {
					try {
						Point2D e = zoom.inverseTransform(event.getPoint(), new Point2D.Double());
						System.out.println(e.getX()+ "  "+e.getY());
						int[] pix = imgI.getSubimage((int)e.getX()-neighbours, (int)e.getY()-neighbours, nn, nn).getRGB(0, 0, nn, nn, new int[nn*nn], 0, nn);
						window.setRGB(0, 0, nn, nn, pix, 0, nn);
						windowShow.repaint();
						double u = 0, v = 0;
						for (int z = 0; z<Z; z++) {
							
							double xx = 0, yy = 0, yx = 0, tx = 0, ty = 0, weight = 0;
							for (int j = -neighbours; j<neighbours; j++) {
								for (int i = -neighbours; i<neighbours; i++) {
									int index = (int)(e.getX()+i+(e.getY()+j)*w), index2 = (int)(e.getX()+u+i+(e.getY()+v+j)*w);
									double gauss = Gauss.gaussian1d(Math.sqrt(i*i+j*j), neighbours/3.);
//									gauss = 1;
//									int dxv = bilerp(dx, x, y, w, h), dyv = bilerp(dy, x, y, w, h), dtv = bilerp(dt, x+u, y+v, w, h);
									int dxv = dx[index], dyv = dy[index], dtv = dt[index2];
//									System.out.println(dxv+" "+dyv+" "+dtv);
									xx += dxv*dxv*gauss;
									yy += dyv*dyv*gauss;
									yx += dyv*dxv*gauss;
									tx -= dtv*dxv*gauss;
									ty -= dtv*dyv*gauss;
									weight += gauss;
								}
							}
							
							xx /= weight;
							yy /= weight;
							yx /= weight;
							tx /= weight;
							ty /= weight;
							
							if (Double.compare(0, xx*yy-yx*yx)!=0) {
								SimpleMatrix sm = new SimpleMatrix(2, 2);
								sm.set(0, 0, xx);
								sm.set(1, 0, yx);
								sm.set(0, 1, yx);
								sm.set(1, 1, yy);
								SimpleEVD<SimpleMatrix> evd = sm.eig();
//								System.out.println(evd.getEigenvalue(0)+" "+evd.getEigenvalue(1));
								if (evd.getEigenvalue(0).real>limit || evd.getEigenvalue(1).real>limit) {
									
									double u_ = ((tx*yy-yx*ty)/(xx*yy-yx*yx));
									double v_ = ((ty-yx*u_)/yy);
									u+=u_;
									v+=v_;
//									if (u*u+v*v<norm) {
//										u *= Math.sqrt(norm);
//										v *= Math.sqrt(norm);
								
									pix = imgJ.getSubimage((int)e.getX()-neighbours+(int)u, (int)e.getY()-neighbours+(int)v, nn, nn).getRGB(0, 0, nn, nn, pix, 0, nn);
									movedwindow.setRGB(0, 0, nn, nn, pix, 0, nn);
								}
							} 
							else {
//								System.out.println("verkackt");
							}

						}
						
								
						movedWindowShow.repaint();
						
					}
					catch (NoninvertibleTransformException nte) {
						nte.printStackTrace();
					}
					super.mouseClicked(event);
				}
			});
			
			JFrame frame = new JFrame();
			frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
			frame.add(something,BorderLayout.CENTER);
			JPanel pan = new JPanel();
			pan.setLayout(new BoxLayout(pan, BoxLayout.Y_AXIS));
			
			pan.add(windowShow);
			pan.add(movedWindowShow);
			
			frame.add(pan,BorderLayout.EAST);
			
			frame.pack();
			frame.setVisible(true);
			

		} catch (Exception e) {	e.printStackTrace(); }
		
		
		
		
		
		
		
		
	} // end main()

	static double subpixel(int[] lum, double col, double row,int w) { // col=spalte, row=zeile, w=weite
		col = Math.min(Math.max(col, 0), w-1);
		row = Math.max(Math.min(row, lum.length/w-1), 0);
		int col0 = (int)Math.floor(col);
		double cola = col-col0;
		int row0 = (int)Math.floor(row);
		double rowa = row-row0;
		double pixel = (1-cola)*(1-rowa)*lum[col0+row0*w] + cola*(1-rowa)*lum[col0+1+row0*w] + (1-cola)*rowa*lum[col0+(row0+1)*w] + cola*rowa*lum[col0+1+(row0+1)*w];
		return pixel;
	}
	
		
	static SimpleMatrix buildLevel(SimpleMatrix L0){
		int nx0 = L0.numCols(), ny0 = L0.numRows();
		int nx1 = (nx0+1)/2, ny1 = (ny0+1)/2; // REF[1] Gleichung (3),(4)
		SimpleMatrix L1 = new SimpleMatrix(ny1,nx1);
		for (int row=0;row<nx1;row++) { for (int col=0;col<ny1;col++) {
			// dummy pixels -- REF[1] §2.1 
			if (row==-1) row=0;	if (col==-1) col=0; // TO-DO :: reicht das aus? was ist mit nx_L-1 und ny_L-1
			// REF[1] Gleichung (2)
			double value = (L0.get(2*row, 2*col))/4. 
				+ (L0.get(2*row-1,2*col)+L0.get(2*row+1,2*col)+L0.get(2*row, 2*col-1)+L0.get(2*row, 2*col+1))/8.
				+ (L0.get(2*row-1,2*col-1)+L0.get(2*row+1,2*col+1)+L0.get(2*row+1, 2*col-1)+L0.get(2*row-1, 2*col+1))/16.;
			L1.set(row, col, value);
		} } // TO-DO :: anti-aliasing filter -- see §2.1
		return L1;
	}
	
		
		
	
	
	
} // end Aufgabe05_w
